from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc

import numpy as np
from tensorflow.python.framework.tensor_util import make_tensor_proto
import grpc
from tensorboard.util.tensor_util import make_ndarray


import os
import requests
import numpy as np
import tensorflow as tf
from nlp_tools.processors.tools import load_processors_from_model

model_save_path = '/home/qiufengfeng/nlp/train_models/classification/'
pd_save_path = os.path.join(model_save_path,'pd_path')
text_processor, label_processor = load_processors_from_model(model_save_path)

text = "这是一个测试"
tensor = np.array(text_processor.transform([text]))


host = "localhost:8500"

#channel = grpc.insecure_channel(host)
import grpc
channel = grpc.insecure_channel("127.0.0.1:8500")


stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

request = predict_pb2.PredictRequest()
request.model_spec.name='classification'
request.model_spec.signature_name="serving_default"

request.inputs["Input-Token"].CopyFrom(make_tensor_proto(tensor.tolist(),dtype=tf.float32)) # shape跟 keras的model.input类型对应
request.inputs["Input-Segment"].CopyFrom(make_tensor_proto(np.zeros(tensor.shape).tolist(),dtype=tf.float32)) # shape跟 keras的model.input类型对应
result = stub.Predict(request, 10.0)  # 10 secs timeout
print(result)

results = {}
for key in result.outputs:
    tensor_proto = result.outputs[key]
    results[key] =  make_ndarray(tensor_proto)
print(results)